
/**
  ******************************************************************************
  * Copyright 2021 The Microbee Authors. All Rights Reserved.
  * 
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  * 
  * http://www.apache.org/licenses/LICENSE-2.0
  * 
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  * 
  * @file       sim_motor.c
  * @author     baiyang
  * @date       2022-11-28
  ******************************************************************************
  */

/*----------------------------------include-----------------------------------*/
#include "sim_motor.h"

#include <stdio.h>
#include <string.h>

#include <common/time/gp_time.h>
/*-----------------------------------macro------------------------------------*/

/*----------------------------------typedef-----------------------------------*/

/*---------------------------------prototype----------------------------------*/

/*----------------------------------variable----------------------------------*/

/*-------------------------------------os-------------------------------------*/

/*----------------------------------function----------------------------------*/
void sim_motor_ctor(sim_motor_t motor, uint8_t _servo, float _angle, float _yaw_factor, uint8_t _display_order)
{
    memset(motor, 0, sizeof(struct sim_motor));

    motor->roll_servo = -1;
    motor->pitch_servo = -1;
    motor->servo_rate = 0.24f; // seconds per 60 degrees

    motor->servo = _servo;      // what servo output drives this motor
    motor->angle = _angle;      // angle in degrees from front
    motor->yaw_factor = _yaw_factor;        // positive is clockwise
    motor->display_order = _display_order;  // order for clockwise display

    motor->position.x = cosf(radians(motor->angle));
    motor->position.y =  sinf(radians(motor->angle));
    motor->position.z = 0.0f;

    motor->thrust_vector.x = 0.0f;
    motor->thrust_vector.y = 0.0f;
    motor->thrust_vector.z = -1.0f;
}

void sim_motor_ctor2(sim_motor_t motor, uint8_t _servo, float _angle, float _yaw_factor, uint8_t _display_order,
          int8_t _roll_servo, float _roll_min, float _roll_max,
          int8_t _pitch_servo, float _pitch_min, float _pitch_max)
{
    memset(motor, 0, sizeof(struct sim_motor));

    motor->roll_servo = -1;
    motor->pitch_servo = -1;
    motor->servo_rate = 0.24; // seconds per 60 degrees

    motor->servo = _servo; // what servo output drives this motor
    motor->angle = _angle; // angle in degrees from front
    motor->yaw_factor = _yaw_factor; // positive is clockwise
    motor->display_order = _display_order; // order for clockwise display
    motor->roll_servo = _roll_servo;
    motor->roll_min = _roll_min;
    motor->roll_max = _roll_max;
    motor->pitch_servo = _pitch_servo;
    motor->pitch_min = _pitch_min;
    motor->pitch_max = _pitch_max;

    motor->position.x = cosf(radians(motor->angle));
    motor->position.y =  sinf(radians(motor->angle));
    motor->position.z = 0.0f;

    motor->thrust_vector.x = 0.0f;
    motor->thrust_vector.y = 0.0f;
    motor->thrust_vector.z = -1.0f;
}

void sim_motor_init(sim_motor_t motor)
{
    motor->position.x = cosf(radians(motor->angle));
    motor->position.y =  sinf(radians(motor->angle));
    motor->position.z = 0.0f;

    motor->thrust_vector.x = 0.0f;
    motor->thrust_vector.y = 0.0f;
    motor->thrust_vector.z = -1.0f;
}

// calculate rotational accel and thrust for a motor
void sim_motor_calculate_forces(sim_motor_t motor, const struct sitl_input *input,
                             uint8_t motor_offset,
                             Vector3f_t *torque,
                             Vector3f_t *thrust,
                             const Vector3f_t *velocity_air_bf,
                             const Vector3f_t *gyro,
                             float air_density,
                             float voltage,
                             bool use_drag)
{

    const float pwm = input->servos[motor_offset+motor->servo];
    float command = sim_motor_pwm_to_command(motor, pwm);
    float voltage_scale = voltage / motor->voltage_max;

    if (voltage_scale < 0.1) {
        // battery is dead
        vec3_zero(torque);
        vec3_zero(thrust);
        motor->current = 0;
        return;
    }

    // apply slew limiter to command
    uint64_t now_us = time_micros64();
    if (motor->last_calc_us != 0 && motor->slew_max > 0) {
        float dt = (now_us - motor->last_calc_us)*1.0e-6;
        float slew_max_change = motor->slew_max * dt;
        command = math_constrain_float(command, motor->last_command-slew_max_change, motor->last_command+slew_max_change);
    }
    motor->last_calc_us = now_us;
    motor->last_command = command;

    // velocity of motor through air
    Vector3f_t motor_vel = *velocity_air_bf;

    // add velocity of motor about center due to vehicle rotation
    //motor_vel += -(motor->position % gyro);
    motor_vel.x += -(motor->position.y*gyro->z - motor->position.z*gyro->y);
    motor_vel.y += -(motor->position.z*gyro->x - motor->position.x*gyro->z);
    motor_vel.z += -(motor->position.x*gyro->y - motor->position.y*gyro->x);

    // calculate velocity into prop, clipping at zero
    float velocity_in = MAX(0, -(vec3_projected(&motor_vel, &motor->thrust_vector).z));

    // get thrust for untilted motor
    float motor_thrust = sim_motor_calc_thrust(motor, command, air_density, velocity_in, voltage_scale);

    // the yaw torque of the motor
    const float yaw_scale = 0.05 * motor->diagonal_size * motor_thrust;
    Vector3f_t rotor_torque = {motor->thrust_vector.x * motor->yaw_factor * command * yaw_scale * -1.0,
                               motor->thrust_vector.y * motor->yaw_factor * command * yaw_scale * -1.0,
                               motor->thrust_vector.z * motor->yaw_factor * command * yaw_scale * -1.0};

    // thrust in bodyframe NED
    thrust->x = motor->thrust_vector.x * motor_thrust;
    thrust->y = motor->thrust_vector.y * motor_thrust;
    thrust->z = motor->thrust_vector.z * motor_thrust;

    // work out roll and pitch of motor relative to it pointing straight up
    float roll = 0, pitch = 0;

    uint64_t now = time_micros64();
    
    // possibly roll and/or pitch the motor
    if (motor->roll_servo >= 0) {
        uint16_t servoval = sim_motor_update_servo(motor, input->servos[motor->roll_servo+motor_offset], now, &motor->last_roll_value);
        if (motor->roll_min < motor->roll_max) {
            roll = math_constrain_float(motor->roll_min + (servoval-1000)*0.001*(motor->roll_max-motor->roll_min), motor->roll_min, motor->roll_max);
        } else {
            roll = math_constrain_float(motor->roll_max + (2000-servoval)*0.001*(motor->roll_min-motor->roll_max), motor->roll_max, motor->roll_min);
        }
    }
    if (motor->pitch_servo >= 0) {
        uint16_t servoval = sim_motor_update_servo(motor, input->servos[motor->pitch_servo+motor_offset], now, &motor->last_pitch_value);
        if (motor->pitch_min < motor->pitch_max) {
            pitch = math_constrain_float(motor->pitch_min + (servoval-1000)*0.001*(motor->pitch_max-motor->pitch_min), motor->pitch_min, motor->pitch_max);
        } else {
            pitch = math_constrain_float(motor->pitch_max + (2000-servoval)*0.001*(motor->pitch_min-motor->pitch_max), motor->pitch_max, motor->pitch_min);
        }
    }
    motor->last_change_usec = now;

    // calculate torque in newton-meters
    torque->x = (motor->position.y*thrust->z - motor->position.z*thrust->y) + rotor_torque.x;
    torque->y = (motor->position.z*thrust->x - motor->position.x*thrust->z) + rotor_torque.y;
    torque->z = (motor->position.x*thrust->y - motor->position.y*thrust->x) + rotor_torque.z;

    // possibly rotate the thrust vector and the rotor torque
    if (!math_flt_zero(roll) || !math_flt_zero(pitch)) {
        matrix3f_t rotation;
        matrix3f_from_euler(&rotation , radians(roll), radians(pitch), 0);
        matrix3f_mul_vec(thrust, &rotation, thrust);
        matrix3f_mul_vec(torque, &rotation, torque);
    }

    if (use_drag) {
        // calculate momentum drag per motor
        const float momentum_drag_factor = motor->momentum_drag_coefficient * sqrtf(air_density * motor->true_prop_area);
        Vector3f_t momentum_drag;
        momentum_drag.x = momentum_drag_factor * motor_vel.x * (sqrtf(fabsf(thrust->y)) + sqrtf(fabsf(thrust->z)));
        momentum_drag.y = momentum_drag_factor * motor_vel.y * (sqrtf(fabsf(thrust->x)) + sqrtf(fabsf(thrust->z)));
        // The application of momentum drag to the Z axis is a 'hack' to compensate for incorrect modelling
        // of the variation of thust with inflow velocity. If not applied, the vehicle will
        // climb at an unrealistic rate during operation in STABILIZE. TODO replace prop and motor model in
        // with one based on DC motor, momentum disc and blade element theory.
        momentum_drag.z = momentum_drag_factor * motor_vel.z * (sqrtf(fabsf(thrust->x)) + sqrtf(fabsf(thrust->y)) + sqrtf(fabsf(thrust->z)));

        thrust->x -= momentum_drag.x;
        thrust->y -= momentum_drag.y;
        thrust->z -= momentum_drag.z;
    }

    // calculate current
    float power = motor->power_factor * fabsf(motor_thrust);
    motor->current = power / MAX(voltage, 0.1);
}

/*
  update and return current value of a servo. Calculated as 1000..2000
 */
uint16_t sim_motor_update_servo(sim_motor_const_t motor, uint16_t demand, uint64_t time_usec, float *last_value)
{
    if (motor->servo_rate <= 0) {
        return demand;
    }
    if (motor->servo_type == SERVO_RETRACT) {
        // handle retract servos
        if (demand > 1700) {
            demand = 2000;
        } else if (demand < 1300) {
            demand = 1000;
        } else {
            demand = *last_value;
        }
    }
    demand = math_constrain_int16(demand, 1000, 2000);
    float dt = (time_usec - motor->last_change_usec) * 1.0e-6f;
    // assume servo moves through 90 degrees over 1000 to 2000
    float max_change = 1000 * (dt / motor->servo_rate) * 60.0f / 90.0f;
    *last_value = math_constrain_float(demand, *last_value-max_change, *last_value+max_change);
    return (uint16_t)(*last_value+0.5f);
}

// setup PWM ranges for this motor
void sim_motor_setup_params(sim_motor_t motor, uint16_t _pwm_min, uint16_t _pwm_max, float _spin_min, float _spin_max, float _expo, float _slew_max,
                         float _diagonal_size, float _power_factor, float _voltage_max, float _effective_prop_area,
                         float _velocity_max, const Vector3f_t* _position, const Vector3f_t* _thrust_vector, float _yaw_factor, 
                         float _true_prop_area, float _momentum_drag_coefficient)
{
    motor->mot_pwm_min = _pwm_min;
    motor->mot_pwm_max = _pwm_max;
    motor->mot_spin_min = _spin_min;
    motor->mot_spin_max = _spin_max;
    motor->mot_expo = _expo;
    motor->slew_max = _slew_max;
    motor->power_factor = _power_factor;
    motor->voltage_max = _voltage_max;
    motor->effective_prop_area = _effective_prop_area;
    motor->max_outflow_velocity = _velocity_max;
    motor->true_prop_area = _true_prop_area;
    motor->momentum_drag_coefficient = _momentum_drag_coefficient;
    motor->diagonal_size = _diagonal_size;

    if (!vec3_is_zero(_position)) {
        motor->position = *_position;
    } else {
        motor->position.x = cosf(radians(motor->angle)) * _diagonal_size;
        motor->position.y =  sinf(radians(motor->angle)) * _diagonal_size;
        motor->position.z = 0;
    }

    if (!vec3_is_zero(_thrust_vector)) {
        motor->thrust_vector = *_thrust_vector;
    }
    if (!math_flt_zero(_yaw_factor)) {
        motor->yaw_factor = _yaw_factor;
    }
}

/*
convert a PWM value to a command value from 0 to 1
*/
float sim_motor_pwm_to_command(sim_motor_const_t motor, float pwm)
{
    const float pwm_thrust_max = motor->mot_pwm_min + motor->mot_spin_max * (motor->mot_pwm_max - motor->mot_pwm_min);
    const float pwm_thrust_min = motor->mot_pwm_min + motor->mot_spin_min * (motor->mot_pwm_max - motor->mot_pwm_min);
    const float pwm_thrust_range = pwm_thrust_max - pwm_thrust_min;
    return math_constrain_float((pwm-pwm_thrust_min)/pwm_thrust_range, 0, 1);
}

/*
  calculate thrust given a command value
*/
float sim_motor_calc_thrust(sim_motor_const_t motor, float command, float air_density, float velocity_in, float voltage_scale)
{
    float velocity_out = voltage_scale * motor->max_outflow_velocity * sqrtf((1-motor->mot_expo)*command + motor->mot_expo*sq(command));
    float ret = 0.5 * air_density * motor->effective_prop_area * (sq(velocity_out) - sq(velocity_in));
#if 0
    if (command > 0) {
        printf("air_density=%f effective_prop_area=%f velocity_in=%f velocity_max=%f\n",
                 air_density, effective_prop_area, velocity_in, voltage_scale * max_outflow_velocity);
        printf("calc_thrust %.3f %.3f\n", command, ret);
    }
#endif
    return ret;
}

/*------------------------------------test------------------------------------*/


